# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from tabulate import tabulate

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 输入数据
data = {
    '日期': pd.date_range(start='2025-08-08', end='2025-09-05'),
    '预测收盘价(元)': [52.46, 52.27, 52.28, 52.29, 52.29, 52.27, 52.25, 52.22,
                       52.19, 52.17, 52.15, 52.13, 52.12, 52.11, 52.10, 52.09,
                       52.08, 52.07, 52.07, 52.06, 52.06, 52.05, 52.04, 52.03,
                       52.03, 52.02, 52.01, 52.01, 52.00],
    '日变化(%)': [np.nan, -0.36, 0.02, 0.03, -0.01, -0.04, -0.05, -0.06,
                  -0.05, -0.05, -0.04, -0.03, -0.02, -0.02, -0.02, -0.02,
                  -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01, -0.01,
                  -0.01, -0.01, -0.01, -0.01, -0.01]
}

df = pd.DataFrame(data)
df['日期'] = pd.to_datetime(df['日期'])

# 策略参数
INITIAL_CAPITAL = 1000000  # 初始资金100万元
FEE_RATE = 0.0003  # 交易费率0.03%
POSITION_RATIO = 0.5  # 仓位系数


def calculate_prices(close_price, scheme):
    """计算买入和卖出价格"""
    low_price = 0.99 * close_price  # 假设最低价为收盘价的99%
    high_price = 1.01 * close_price  # 假设最高价为收盘价的101%

    if scheme == 1:
        return low_price, high_price  # 方案1：买入=最低价，卖出=最高价
    elif scheme == 2:
        # 方案2：买入=最低价+1/4价差，卖出=最低价+3/4价差
        spread = high_price - low_price
        buy_price = low_price + 0.25 * spread
        sell_price = low_price + 0.75 * spread
        return buy_price, sell_price
    else:
        raise ValueError("无效的价格方案")


def simulate_trading(data, scheme):
    """模拟交易过程"""
    cash = INITIAL_CAPITAL
    shares_held = 0
    shares_bought_today = 0  # T+1规则跟踪
    results = []

    for i in range(len(data) - 1):
        current_close = data.loc[i, '预测收盘价(元)']
        delta_p = data.loc[i + 1, '日变化(%)'] / 100  # 转换为小数

        # 计算买卖价格
        buy_price, sell_price = calculate_prices(current_close, scheme)

        # 前一日的买入股票今日可以卖出
        shares_available = shares_held - shares_bought_today

        # 计算交易数量
        if not np.isnan(delta_p):
            trade_value = POSITION_RATIO * abs(delta_p) * (cash + shares_held * current_close)
            trade_shares = trade_value / current_close
        else:
            trade_shares = 0

        # 执行交易
        shares_bought_today = 0  # 重置当日买入量

        if delta_p > 0 and cash > 0:  # 预测上涨，买入
            max_buy_shares = cash / (buy_price * (1 + FEE_RATE))
            actual_shares = min(trade_shares, max_buy_shares)
            cost = actual_shares * buy_price * (1 + FEE_RATE)
            cash -= cost
            shares_held += actual_shares
            shares_bought_today = actual_shares
            action = f"买入{actual_shares:.2f}股"
        elif delta_p < 0 and shares_available > 0:  # 预测下跌，卖出
            actual_shares = min(trade_shares, shares_available)
            proceeds = actual_shares * sell_price * (1 - FEE_RATE)
            cash += proceeds
            shares_held -= actual_shares
            action = f"卖出{actual_shares:.2f}股"
        else:
            action = "持有"

        # 计算总资产
        total_assets = cash + shares_held * current_close

        # 记录结果
        results.append({
            '日期': data.loc[i, '日期'].strftime('%Y-%m-%d'),
            '收盘价': current_close,
            '预测变化(%)': delta_p * 100 if not np.isnan(delta_p) else np.nan,
            '操作': action,
            '买入价': buy_price if '买入' in action else np.nan,
            '卖出价': sell_price if '卖出' in action else np.nan,
            '持股数量': shares_held,
            '现金余额': cash,
            '总资产': total_assets
        })

    return pd.DataFrame(results)


def analyze_results(results, scheme_name):
    """分析策略结果"""
    initial = INITIAL_CAPITAL
    final = results.iloc[-1]['总资产']
    returns = (final - initial) / initial * 100

    transactions = results[results['操作'] != '持有']

    # 更稳健的费用计算方式
    buy_trans = transactions[transactions['操作'].str.contains('买入')]
    sell_trans = transactions[transactions['操作'].str.contains('卖出')]

    # 提取交易数量
    buy_shares = buy_trans['操作'].str.extract(r'买入(\d+\.\d+)')[0].astype(float)
    sell_shares = sell_trans['操作'].str.extract(r'卖出(\d+\.\d+)')[0].astype(float)

    # 计算总费用
    buy_fees = (buy_trans['买入价'] * buy_shares * FEE_RATE).sum()
    sell_fees = (sell_trans['卖出价'] * sell_shares * FEE_RATE).sum()
    total_fees = buy_fees + sell_fees

    print(f"\n{scheme_name}策略结果分析:")
    print(f"- 初始资产: {initial:,.2f}元")
    print(f"- 最终资产: {final:,.2f}元")
    print(f"- 总收益率: {returns:.2f}%")
    print(f"- 交易次数: {len(transactions)}次")
    print(f"- 总交易费用: {total_fees:,.2f}元")

    # 绘制资产变化曲线
    plt.figure(figsize=(12, 6))
    plt.plot(pd.to_datetime(results['日期']), results['总资产'], label='总资产')
    plt.title(f'{scheme_name}策略资产变化')
    plt.xlabel('日期')
    plt.ylabel('资产总额(元)')
    plt.grid(True)
    plt.legend()
    plt.show()

    return results


# 方案1模拟
print("=" * 50)
print("方案1：卖出=当日最高价，买入=当日最低价")
results_scheme1 = simulate_trading(df, scheme=1)
analyze_results(results_scheme1, "方案1")

# 方案2模拟
print("=" * 50)
print("方案2：卖出=3/4分位点，买入=1/4分位点")
results_scheme2 = simulate_trading(df, scheme=2)
analyze_results(results_scheme2, "方案2")

# 对比两种方案
print("=" * 50)
print("两种方案对比:")
final1 = results_scheme1.iloc[-1]['总资产']
final2 = results_scheme2.iloc[-1]['总资产']
print(f"- 方案1最终资产: {final1:,.2f}元")
print(f"- 方案2最终资产: {final2:,.2f}元")
print(f"- 方案1比方案2多盈利: {final1 - final2:,.2f}元 ({(final1 - final2) / INITIAL_CAPITAL * 100:.2f}%)")

# 输出详细交易记录
print("\n方案1前5日交易记录:")
print(tabulate(results_scheme1.head(), headers='keys', tablefmt='grid', showindex=False))

print("\n方案2前5日交易记录:")
print(tabulate(results_scheme2.head(), headers='keys', tablefmt='grid', showindex=False))